﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using PerfAssist.Common.usmooth.Common;
using PerfAssist.Common.usmooth.Runtime;
using UnityEngine;
using UnityEngine.SceneManagement;
using Debug = UnityEngine.Debug;
using Object = UnityEngine.Object;
using Random = UnityEngine.Random;

namespace PerfAssist.MemoryProfilerAdvanced
{
    public enum ResourceRequestType
    {
        Ordinary,
        Async
    }

    public class ResourceRequestInfo
    {
        public double requestTime = Time.realtimeSinceStartup;

        public ResourceRequestType requestType;

        public string resourcePath = "";
        public Type resourceType;
        public int rootID;
        public int seqID;

        public string srcFile = "";
        public int srcLineNum;

        public int stacktraceHash;

        public override string ToString()
        {
            return string.Format("#{0} ({1:0.000}) {2} {3} {4} +{5} +{6} ({7})",
                seqID, requestTime, rootID, resourceType != null ? resourceType.ToString() : null,
                requestType == ResourceRequestType.Async ? "(a)" : "", resourcePath, srcFile, srcLineNum);
        }

        public void RecordObject(Object obj)
        {
            rootID = obj.GetInstanceID();
            resourceType = obj.GetType();
        }
    }

    public class ResourceTracker : IDisposable
    {
        public static ResourceTracker Instance;

        private StreamWriter _logWriter;

        //private string _logPath = "";
        private int _reqSeq;

        private readonly Dictionary<string, List<stackParamater>> _stackUnavailableDict =
            new Dictionary<string, List<stackParamater>>();

        private readonly Dictionary<object, ResourceRequestInfo> InProgressAsyncObjects =
            new Dictionary<object, ResourceRequestInfo>();

        private readonly Dictionary<int, string> Stacktraces = new Dictionary<int, string>();
        private readonly Dictionary<int, ResourceRequestInfo> TrackedAllocInfo = new Dictionary<int, ResourceRequestInfo>();

        private readonly Dictionary<int, int> TrackedGameObjects = new Dictionary<int, int>();
        private readonly Dictionary<int, int> TrackedMemObjects = new Dictionary<int, int>();

        public ResourceTracker(bool enableTracking)
        {
            if (enableTracking)
            {
                Open();

                if (EnableTracking)
                {
                    if (UsNet.Instance != null && UsNet.Instance.CmdExecutor != null)
                    {
                        UsNet.Instance.CmdExecutor.RegisterHandler(eNetCmd.CL_RequestStackData, NetHandle_RequestStackData);
                        UsNet.Instance.CmdExecutor.RegisterHandler(eNetCmd.CL_RequestStackSummary,
                            NetHandle_RequestStackSummary);
                    }
                    else
                    {
                        Debug.LogError("UsNet not available");
                    }

                    readShaderPropertyJson();
                }
            }
        }

        // this boolean is *read-only* after the instance is created
        public bool EnableTracking { get; private set; }

        public Dictionary<string, string> ShaderPropertyDict { get; private set; }

        public void Dispose()
        {
            if (EnableTracking && _logWriter != null) Close();
        }

        private void readShaderPropertyJson()
        {
            if (ShaderPropertyDict == null)
                try
                {
                    var sr = new StreamReader(
                        new FileStream(ResourceTrackerConst.shaderPropertyNameJsonPath, FileMode.Open));
                    var jsonStr = sr.ReadToEnd();
                    sr.Close();
                    ShaderPropertyDict = JsonUtility.FromJson<Serialization<string, string>>(jsonStr).ToDictionary();
                }
                catch (Exception)
                {
                    Debug.Log("no ShaderPropertyNameRecord.json");
                }
        }


        public void Open()
        {
            if (EnableTracking)
            {
                Debug.LogFormat("[ResourceTracker] info: {0} ", "already enabled, ignored.");
                return;
            }

            try
            {
                var logDir = Path.Combine(Application.persistentDataPath, "mem_logs");
                if (!Directory.Exists(logDir))
                    Directory.CreateDirectory(logDir);

                var dt = DateTime.Now;

                var logFile = string.Format("{0}_{1}_alloc.txt", SysUtil.FormatDateAsFileNameString(dt),
                    SysUtil.FormatTimeAsFileNameString(dt));
                var logPath = Path.Combine(logDir, logFile);

                _logWriter = new FileInfo(logPath).CreateText();
                _logWriter.AutoFlush = true;
                //_logPath = logPath;

                EnableTracking = true;
                Debug.LogFormat("[ResourceTracker] tracking enabled: {0} ", logPath);
            }
            catch (Exception ex)
            {
                Debug.LogErrorFormat("[ResourceTracker] Open() failed, error: {0} ", ex.Message);

                if (_logWriter != null)
                {
                    _logWriter.Close();
                    _logWriter = null;
                }

                EnableTracking = false;
                //_logPath = "";
            }
        }

        public void Close()
        {
            if (_logWriter != null)
                try
                {
                    _logWriter.WriteLine("--------- unfinished request: {0} --------- ", InProgressAsyncObjects.Count);
                    foreach (var p in InProgressAsyncObjects) _logWriter.WriteLine("  + {0}", p.Value);
                    _logWriter.Close();
                }
                catch (Exception ex)
                {
                    Debug.LogErrorFormat("[ResourceTracker.ctor] error: {0} ", ex.Message);
                }
                finally
                {
                    _logWriter = null;
                }

            EnableTracking = false;
        }

        public void TrackSyncRequest(Object spawned, string path)
        {
            if (!EnableTracking)
                return;

            var sf = new StackFrame(2, true);
            var request = NewRequest(path, sf);
            request.requestType = ResourceRequestType.Ordinary;
            TrackRequestWithObject(request, spawned);
        }

        public void TrackResourcesDotLoad(Object loaded, string path)
        {
            if (!EnableTracking)
                return;

            var sf = new StackFrame(1, true);
            var request = NewRequest(path, sf);
            request.requestType = ResourceRequestType.Ordinary;
            TrackRequestWithObject(request, loaded);
        }

        public void TrackAsyncRequest(object handle, string path)
        {
            if (!EnableTracking)
                return;

            var sf = new StackFrame(2, true);
            if (sf.GetMethod().Name.Contains("SpawnAsyncOldVer")) sf = new StackFrame(3, true);

            InProgressAsyncObjects[handle] = NewRequest(path, sf);
        }

        public void TrackAsyncDone(object handle, Object target)
        {
            if (!EnableTracking || target == null)
                return;

            ResourceRequestInfo request;
            if (!InProgressAsyncObjects.TryGetValue(handle, out request))
                return;

            request.requestType = ResourceRequestType.Async;
            TrackRequestWithObject(request, target);
            InProgressAsyncObjects.Remove(handle);
        }

        public void TrackSceneLoaded(string sceneName)
        {
            if (!EnableTracking)
                return;

            var scene = SceneManager.GetActiveScene();
            var roots = scene.GetRootGameObjects();
            for (var i = 0; i < roots.Length; i++) TrackSyncRequest(roots[i], "[scene]: " + sceneName);
        }

        public void TrackObjectInstantiation(Object src, Object instantiated)
        {
            if (!EnableTracking)
                return;

            var allocSeqID = -1;
            if (!TrackedGameObjects.TryGetValue(src.GetInstanceID(), out allocSeqID))
                return;

            ExtractObjectResources(instantiated, allocSeqID);
        }

        public ResourceRequestInfo GetAllocInfo(int instID)
        {
            if (!EnableTracking)
                return null;

            var allocSeqID = -1;

            if (!TrackedGameObjects.TryGetValue(instID, out allocSeqID) &&
                !TrackedMemObjects.TryGetValue(instID, out allocSeqID))
                return null;

            ResourceRequestInfo requestInfo = null;
            if (!TrackedAllocInfo.TryGetValue(allocSeqID, out requestInfo))
                return null;

            return requestInfo;
        }

        public string GetStackTrace(ResourceRequestInfo req)
        {
            string stacktrace;
            if (!Stacktraces.TryGetValue(req.stacktraceHash, out stacktrace))
                return "";

            return stacktrace;
        }

        private ResourceRequestInfo NewRequest(string path, StackFrame sf)
        {
            var reqInfo = new ResourceRequestInfo();
            reqInfo.resourcePath = path;
            reqInfo.srcFile = sf.GetFileName();
            reqInfo.srcLineNum = sf.GetFileLineNumber();
            reqInfo.seqID = _reqSeq++;

            var stacktrace = StackTraceUtility.ExtractStackTrace();

            var _tryCount = 10;
            while (_tryCount > 0)
            {
                string stacktraceStored;
                if (!Stacktraces.TryGetValue(stacktrace.GetHashCode(), out stacktraceStored))
                {
                    Stacktraces[stacktrace.GetHashCode()] = stacktrace;
                    break;
                }

                if (stacktrace == stacktraceStored)
                    break;
                stacktrace += ((int) (Random.value * 100)).ToString();

                _tryCount--;
            }

            reqInfo.stacktraceHash = stacktrace.GetHashCode();
            return reqInfo;
        }

        private void ExtractObjectResources(Object obj, int reqSeqID)
        {
            var sge = new SceneGraphExtractor(obj);

            for (var i = 0; i < sge.GameObjectIDs.Count; i++)
                if (!TrackedGameObjects.ContainsKey(sge.GameObjectIDs[i]))
                    TrackedGameObjects[sge.GameObjectIDs[i]] = reqSeqID;

            foreach (var p in sge.MemObjectIDs)
            foreach (var item in p.Value)
                if (!TrackedMemObjects.ContainsKey(item))
                    TrackedMemObjects[item] = reqSeqID;
        }

        public bool NetHandle_RequestStackSummary(eNetCmd cmd, UsCmd c)
        {
            var flag = c.ReadString();
            if (string.IsNullOrEmpty(flag))
                return false;

            if (flag.Equals("begin"))
            {
                _stackUnavailableDict.Clear();
                return true;
            }

            if (flag.Equals("end"))
            {
                Debug.Log("Stack Category Statistical Information:");
                //NetUtil.Log("堆栈类型统计信息:");
                var totalCount = 0;
                var unavailableTotalCount = 0;
                var totalSize = 0;
                var unavailableTotalSize = 0;
                var categoryCount = c.ReadInt32();
                for (var i = 0; i < categoryCount; i++)
                {
                    var category = c.ReadString();
                    List<stackParamater> unavailableList;
                    _stackUnavailableDict.TryGetValue(category, out unavailableList);
                    if (unavailableList != null)
                    {
                        var CategoryCount = c.ReadInt32();
                        var CategorySize = c.ReadInt32();
                        totalCount += CategoryCount;
                        totalSize += CategorySize;
                        unavailableTotalCount += unavailableList.Count;
                        var categoryTotalSize = 0;
                        foreach (var info in unavailableList) categoryTotalSize += info.Size;
                        unavailableTotalSize += categoryTotalSize;
                        Debug.Log(string.Format("[{0} =({1}/{2},{3}/{4})]", category, unavailableList.Count, CategoryCount,
                            ResourceTrackerConst.FormatBytes(categoryTotalSize),
                            ResourceTrackerConst.FormatBytes(CategorySize)));
                        //NetUtil.Log("【{0} =({1}/{2},{3}/{4})】", category, unavailableList.Count, CategoryCount, ResourceTrackerConst.FormatBytes(categoryTotalSize), ResourceTrackerConst.FormatBytes(CategorySize));
                    }
                }

                Debug.Log(string.Format("[total =({0}/{1},{2}/{3})]", unavailableTotalCount, totalCount,
                    ResourceTrackerConst.FormatBytes(unavailableTotalSize), ResourceTrackerConst.FormatBytes(totalSize)));
                //NetUtil.Log("【total =({0}/{1},{2}/{3})】", unavailableTotalCount, totalCount, ResourceTrackerConst.FormatBytes(unavailableTotalSize), ResourceTrackerConst.FormatBytes(totalSize));
                return true;
            }

            var className = flag;
            var count = c.ReadInt32();
            for (var i = 0; i < count; i++)
            {
                var instanceID = c.ReadInt32();
                var size = c.ReadInt32();
                var requestInfo = Instance.GetAllocInfo(instanceID);
                if (requestInfo == null)
                {
                    if (!_stackUnavailableDict.ContainsKey(className))
                        _stackUnavailableDict.Add(className, new List<stackParamater>());
                    List<stackParamater> stackUnavailableList;
                    _stackUnavailableDict.TryGetValue(className, out stackUnavailableList);
                    var info = new stackParamater();
                    info.InstanceID = instanceID;
                    info.Size = size;
                    stackUnavailableList.Add(info);
                }
            }

            return true;
        }

        public List<Texture2D> GetTexture2DObjsFromMaterial(Material mat)
        {
            var shaderPropertyDict = Instance.ShaderPropertyDict;
            var result = new List<Texture2D>();
            if (mat != null)
            {
                var shader = mat.shader;
                if (shader != null && shaderPropertyDict != null && shaderPropertyDict.ContainsKey(shader.name))
                {
                    string propertyNameStrs;
                    shaderPropertyDict.TryGetValue(shader.name, out propertyNameStrs);
                    char[] tokens = {ResourceTrackerConst.shaderPropertyNameJsonToken};
                    var propertyNameList = propertyNameStrs.Split(tokens);
                    foreach (var propertyName in propertyNameList)
                    {
                        var tex = mat.GetTexture(propertyName) as Texture2D;
                        if (tex != null)
                            result.Add(tex);
                        //UnityEngine.Debug.Log(string.Format("catch UIWidget shader texture2D ,textureName = {0}", tex.name));
                    }
                }

                var mainTexture2D = mat.mainTexture as Texture2D;
                if (mainTexture2D != null && !result.Contains(mainTexture2D))
                    result.Add(mainTexture2D);
                //UnityEngine.Debug.Log(string.Format("catch UIWidget shader texture2D ,textureName = {0}", mat.mainTexture.name));
            }

            return result;
        }

        public bool NetHandle_RequestStackData(eNetCmd cmd, UsCmd c)
        {
            var instanceID = c.ReadInt32();
            var className = c.ReadString();
            Debug.Log(string.Format("NetHandle_RequestStackData instanceID={0} className={1}", instanceID, className));

            var requestInfo = Instance.GetAllocInfo(instanceID);

            var pkt = new UsCmd();
            pkt.WriteNetCmd(eNetCmd.SV_QueryStacksResponse);
            if (requestInfo == null)
                pkt.WriteString("<no_callstack_available>");
            else
                pkt.WriteString(Instance.GetStackTrace(requestInfo));
            UsNet.Instance.SendCommand(pkt);
            return true;
        }


        private void TrackRequestWithObject(ResourceRequestInfo req, Object obj)
        {
            if (obj == null)
                return;

            try
            {
                req.RecordObject(obj);

                TrackedAllocInfo[req.seqID] = req;
                ExtractObjectResources(obj, req.seqID);

                if (_logWriter != null)
                    _logWriter.WriteLine(req.ToString());
            }
            catch (Exception ex)
            {
                Debug.LogErrorFormat("[ResourceTracker.TrackRequestWithObject] error: {0} \n {1} \n {2}",
                    ex.Message, req != null ? req.ToString() : "", ex.StackTrace);
            }
        }

        private class stackParamater
        {
            public int InstanceID { get; set; }

            public int Size { get; set; }
        }
    }
}